import numpy as np
import matplotlib.pyplot as plt

# 参数设置
num_followers = 5  # Follower 无人机数量
dt = 0.1           # 时间步长
steps = 200        # 仿真步数

# 增益参数
k_p = 1.2          # 位置控制增益
k_v = 0.7          # 速度控制增益

# 初始化 Leader 状态（确保数据类型为 float）
p_leader = np.array([0.0, 0.0], dtype=float)     # Leader 初始位置
v_leader = np.array([0.15, 0.1], dtype=float)    # Leader 初始速度

# 初始化 Followers 状态
followers = {
    "positions": np.random.rand(num_followers, 2) * 5,  # 初始位置在较小范围内
    "velocities": np.zeros((num_followers, 2)),         # 初始速度为零
    "desired_offsets": np.array([[2 * i, i] for i in range(num_followers)], dtype=float)  # 队形偏移（动态队形）
}

# 仿真循环
plt.figure(figsize=(8, 8))  # 设置更大的画布
for step in range(steps):
    plt.clf()

    # 动态调整队形（模拟队形变化）
    if step % 50 == 0:
        followers["desired_offsets"] = np.array([[2 * i, (-1)**i * i] for i in range(num_followers)], dtype=float)

    # 绘制 Leader
    plt.scatter(p_leader[0], p_leader[1], c='r', label="Leader")

    # 更新 Followers 状态
    for i in range(num_followers):
        p_i = followers["positions"][i]
        v_i = followers["velocities"][i]
        p_desired = p_leader + followers["desired_offsets"][i]  # 计算目标位置

        # 计算控制输入
        u_p = k_p * (p_desired - p_i)  # 位置控制输入
        u_v = k_v * (v_leader - v_i)  # 速度控制输入
        u_i = u_p + u_v               # 综合控制输入

        # 更新位置和速度
        followers["velocities"][i] += u_i * dt
        followers["positions"][i] += followers["velocities"][i] * dt

        # 绘制 Follower 及其与目标点的连接线
        plt.scatter(p_i[0], p_i[1], c='b', label="Follower" if i == 0 else "")  # 避免重复图例
        plt.plot([p_i[0], p_desired[0]], [p_i[1], p_desired[1]], 'g--', linewidth=0.5)

    # 更新 Leader 位置
    p_leader += v_leader * dt

    # 绘制界面
    plt.xlim(-5, 20)
    plt.ylim(-5, 20)
    plt.grid(True)
    plt.title(f"Step: {step}")
    plt.legend(loc="upper left")
    plt.pause(0.1)

plt.show()
